# -*- coding: utf-8 -*-
import logging
import numpy as np
import random
from gym import spaces
import gym
import random
import copy
import Config


#描述了棋盘的相关环境
class deliveryManEnv(gym.Env):

    logger = logging.getLogger(__name__)

    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 2
    }



    #初始化一个n*n的棋盘
    def __init__(self,barrier_states,reward_states,play1state,play2state,
                 play1score,play2score,final1Score,final2Score,play1jobs,play2jobs,
                 play1base=65,play2base=78):

        self.CHECKERBOARDSIZE = Config.CHECKERBOARDSIZE

        self.states = list(range(Config.CHECKERBOARDSIZE * Config.CHECKERBOARDSIZE)) # 状态空间

        self.barrier_states = barrier_states
        self.reward_states  = reward_states
        self.init_reward_states = self.reward_states.copy()

        self.actions = ['U', 'L', 'D', 'R','S']

        self.gamma = Config.GAMMA  # 折扣因子
        self.viewer = None



        # self.action_space = self.getAllLegalMove(self.state)
        self.observation_space = list(set(self.states) - self.barrier_states)

        self.play1state = play1state
        self.play2state = play2state
        self.state = 0

        #复制原版数据
        self.init_state = self.state
        self.init_play1state = play1state
        self.init_play2state = play2state
        self.play1score = play1score
        self.play2score = play2score
        self.init_play1score = play1score
        self.init_play2score = play2score

        #给出最终得分，即送到大本营之后的
        self.final1Score = final1Score
        self.final2Score = final2Score
        self.init_final1Score = final1Score
        self.init_final2Score = final2Score

        #设置大本营
        self.play1base = play1base
        self.play2base = play2base

        #设置背包数量
        self.play1capacity = play1jobs
        self.play2capacity = play2jobs
        self.init_play1capacity = play1jobs
        self.init_play2capacity = play2jobs




    #给出前后左右之后棋盘的取值
    def getActionNum(self,state,action):
        if action == "U":
            result = state - Config.CHECKERBOARDSIZE
        if action == "D":
            result = state + Config.CHECKERBOARDSIZE
        if action == "L":
            result = state - 1
        if action == "R":
            result = state + 1
        if action == "S":
            result = state

        return result

    #判断在当前局面下，走步是否合法
    #判断下一步是不是
    def judgeActionLegal(self,state,action,player):

        newState = self.getActionNum(state,action)

        #超过上下边界
        if newState < 0 or newState > Config.CHECKERBOARDSIZE * Config.CHECKERBOARDSIZE -1:
            return False

        if state % Config.CHECKERBOARDSIZE == 0 and action == "L":
            return False

        if state % Config.CHECKERBOARDSIZE == Config.CHECKERBOARDSIZE - 1 and action == "R":
            return False

        if newState in self.barrier_states:
            return False

        # 不能和敌方或者敌方大本营重合
        if player == 1:
            if newState == self.play2base:
                return  False
        if player == 2:
            if newState == self.play1base:
                return  False

        return True

    def getAllLegalMove(self,state,player):

        resultMove = []
        for move in self.actions:
            if self.judgeActionLegal(state,move,player) == True:
                resultMove.append(move)

        return resultMove


    #给出周围几个点
    def getAllLegalMoveAroud(self,state,player,deepth):

        stateList = set()
        visitedNode = set()
        stateList.add(state)
        for deep in range(deepth):
            toadd = set()
            for nowstate in stateList:
                if nowstate in visitedNode:
                    continue
                moves = self.getAllLegalMove(nowstate,player)
                if "S" in  moves:
                    moves.remove("S")
                for move in moves:
                    newstate = self.getActionNum(nowstate,move)
                    toadd.add(newstate)
                visitedNode.add(nowstate)

            #取并集
            stateList = stateList | toadd

        return stateList




    #每一轮的更新过程,每走一步会衰减reward
    def updateReward(self,next_state):

        rewardKeys = self.reward_states.keys()
        toDelKeys = []

        for key in rewardKeys:
            #如果拿到这个快件，则直接清除
            if key == next_state:
                toDelKeys.append(key)
                continue

            self.reward_states[key] = self.reward_states[key] - self.gamma

            #如果等于0 ，价值为0的 也清除
            if self.reward_states[key] < 0:
                toDelKeys.append(key)

        for delkey in toDelKeys:
            self.reward_states.pop(delkey)


    def getTerminal(self):
        return self.terminate_states

    def getGamma(self):
        return self.gamma

    def getStates(self):
        return self.states

    def getAction(self):
        return self.actions


    def getTerminate_states(self):
        return self.terminate_states


    def setAction(self,s):
        self.state=s

    def step(self, action):
        
        #强制要求合法
        assert self.judgeActionLegal(self.state,action)

        #系统当前状态
        state = self.state

        is_terminal = False

        next_state = self.getActionNum(state,action)
        self.next_state = next_state

        if next_state not in self.reward_states.keys():
            r = 0.0
        else:
            r = self.reward_states[next_state]

        #reward的变化
        self.updateReward(next_state)

        #判断当前是否还有豆，如果没豆，则认为是退出,这个需要验证
        if len(self.reward_states.keys()) <= 0:
            is_terminal = True

        #重置当前节点
        self.state = next_state
        self.action_space = self.getAllLegalMove(self.state)
        return next_state, r,is_terminal,{}

    #带有player的一步
    def playerstep(self, action,player):

        if player == 1:
            # 系统当前状态
            state = self.play1state
        # 系统当前状态
        if player == 2:
            state = self.play2state

        # 强制要求合法
        if (self.judgeActionLegal(state, action,player) == False):
            print(state)
            print(action)
            print(self.barrier_states)
            assert self.judgeActionLegal(state, action)


        is_terminal = False

        next_state = self.getActionNum(state, action)
        if next_state not in self.reward_states.keys():
            r = 0.0
        else:
            r = self.reward_states[next_state]

        # reward的变化
        # self.updateReward(next_state)
        # 判断当前是否还有豆，如果没豆，则认为是退出,这个需要验证


        # 重置当前节点
        if player == 1:
            # 系统当前状态
            self.play1state = next_state
            if self.play1capacity < Config.CAPACITYLIMIT:
                self.play1score = self.play1score + r
                if r > 0:
                    self.play1capacity = self.play1capacity + 1
                self.updateReward(next_state)

            #如果包满了
            else:
                r = 0


            #如果返回大本营  算最终得分
            if next_state == self.play1base:
                self.final1Score = self.final1Score + self.play1score
                self.play1score = 0
                self.play1capacity = 0
                self.updateReward(next_state)

        # 系统当前状态
        if player == 2:
            self.play2state = next_state
            if self.play2capacity < Config.CAPACITYLIMIT:
                self.play2score = self.play2score + r
                if r > 0:
                    self.play2capacity = self.play2capacity + 1

            #包满了 回报就是0了
            else:
                r = 0

            #判断回城
            if next_state == self.play2base:
                self.final2Score = self.final2Score + self.play2score
                self.play2score = 0
                self.play2capacity = 0



        if len(self.reward_states.keys()) <= 0:
            is_terminal = True

        self.action_space = self.getAllLegalMove(self.state,player)
        return next_state, r, is_terminal, {}




    #回复到本次最初的状态
    def reset(self):
        self.state = self.init_state
        self.reward_states = self.init_reward_states.copy()
        # self.action_space = self.getAllLegalMove(self.state)
        self.play1state = self.init_play1state
        self.play2state = self.init_play2state
        self.play1score = self.init_play1score
        self.play2score = self.init_play2score
        self.play1capacity = self.init_play1capacity
        self.play2capacity = self.init_play2capacity
        self.final1Score = self.init_final1Score
        self.final2Score = self.init_final2Score

        return self.state



    #给出一些工具性的函数
    #给出向量和方向转换的函数
    def actionToArray(self,direction):
        if direction == "L":
            return np.array([1,0,0,0,0])
        if direction == "R":
            return np.array([0,1,0,0,0])
        if direction == "D":
            return np.array([0,0,1,0,0])
        if direction == "U":
            return np.array([0,0,0,1,0])
        if direction == "S":
            return np.array([0,0,0,0,1])

    def arrayToAction(self,array):
        num = np.argmax(array)
        if num == 0:
            return "L"
        if num == 1:
            return "R"
        if num == 2:
            return "D"
        if num == 3:
            return "U"
        if num == 4:
            return "S"

    def toAstarBoard(self,player):

        # 定义一个Array 2D类型
        class Array2D:
            """
            说明：
                1、构造方法需要两个参数，即二维数组的宽和高
                2、成员变量w和h是二维数组的宽和高
                3、使用：对象[x][y]可以直接取到相应的值
                4、数组的默认值都是0

            """

            def __init__(self, w, h):
                self.w = w
                self.h = h
                self.data = []
                self.data = [[0 for y in range(h)] for x in range(w)]

            def showarray2d(self):
                for y in range(self.h):
                    for x in range(self.w):
                        print(self.data[x][y], end=' ')

                    print('')

            def __getitem__(self, item):
                return self.data[item]

        map2d = Array2D(12, 12)

        for num in range(0,Config.CHECKERBOARDSIZE * Config.CHECKERBOARDSIZE):
            x = int(num % Config.CHECKERBOARDSIZE)
            y = int(num / Config.CHECKERBOARDSIZE)

            if num in self.barrier_states:
                map2d[x][y] = 1

            if player == 1:
                if num == self.play2base:
                    map2d[x][y] = 1

            if player == 2:
                if num == self.play1base:
                    map2d[x][y] = 1


        return map2d



    def toProgrammingBoard(self,player):

        # 定义一个Array 2D类型

        result = np.zeros([Config.CHECKERBOARDSIZE,Config.CHECKERBOARDSIZE])

        for num in range(0,Config.CHECKERBOARDSIZE * Config.CHECKERBOARDSIZE):
            x = int(num / Config.CHECKERBOARDSIZE)
            y = int(num % Config.CHECKERBOARDSIZE)

            if num in self.barrier_states:
                result[x][y] = 1

            if num in self.reward_states:
                result[x][y] = self.reward_states[num]

            if player == 1:
                if num == self.play2base:
                    result[x][y] = 1

                if num == self.play1state:
                    result[x][y] = 2


            if player == 2:
                if num == self.play1base :
                    result[x][y] = 1

                if num == self.play2state:
                    result[x][y] = 2


        return result

















